package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"regexp"
	"strings"

	"github.com/alecthomas/chroma/v2/formatters"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/styles"
)

// 默认配置
const (
	defaultAPIURL      = "http://localhost:1234/v1/chat/completions"
	defaultModel       = "google/gemma-3-4b"
	defaultTemperature = 0.7
	defaultMaxTokens   = -1
)

// ChatRequest 定义 Ollama API 的请求结构
type ChatRequest struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Temperature float64   `json:"temperature"`
	MaxTokens   int       `json:"max_tokens"`
	Stream      bool      `json:"stream"`
}

// Message 定义消息结构
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// ChatResponse 定义 Ollama API 的响应结构
type ChatResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Error string `json:"error,omitempty"`
}

// 代码块正则表达式
// var codeBlockRegex = regexp.MustCompile(`(?s)(```|~~~)\s*(\w*)\n(.*?)\n\s*(\1)`)
var codeBlockRegex = regexp.MustCompile("(?s)(```|~~~)\\s*(\\w*)\\n(.*?)\\n\\s*(```|~~~)")

// formatCode 格式化代码块，应用语法高亮
func formatCode(content string) string {
	matches := codeBlockRegex.FindAllStringSubmatch(content, -1)
	if len(matches) == 0 {
		return content // 没有代码块，直接返回
	}

	result := content
	for _, match := range matches {
		lang := match[2] // 语言（如 python、go）
		code := match[3] // 代码内容
		if lang == "" {
			lang = "text" // 默认纯文本
		}

		lexer := lexers.Get(lang)
		if lexer == nil {
			lexer = lexers.Fallback
		}
		iterator, err := lexer.Tokenise(nil, code)
		if err != nil {
			return content // 解析失败，返回原内容
		}

		// 使用 One Dark 配色方案（时尚、对比度高）
		formatter := formatters.Get("terminal256")
		if formatter == nil {
			formatter = formatters.Fallback
		}
		var buf bytes.Buffer
		err = formatter.Format(&buf, styles.Get("onedark"), iterator)
		if err != nil {
			return content // 格式化失败，返回原内容
		}

		// 替换原始代码块
		highlighted := buf.String()
		result = strings.Replace(result, match[0], highlighted, 1)
	}
	return result
}

// queryOllama 向 Ollama API 发送请求
func queryOllama(prompt, model string, temperature float64, maxTokens int) (string, error) {
	// 构造请求 payload
	reqBody := ChatRequest{
		Model: model,
		Messages: []Message{
			{Role: "user", Content: prompt},
		},
		Temperature: temperature,
		MaxTokens:   maxTokens,
		Stream:      false,
	}

	// 序列化为 JSON
	body, err := json.Marshal(reqBody)
	if err != nil {
		return "", fmt.Errorf("failed to marshal request: %v", err)
	}

	// 发送 HTTP 请求
	resp, err := http.Post(defaultAPIURL, "application/json", bytes.NewBuffer(body))
	if err != nil {
		return "", fmt.Errorf("failed to send request: %v", err)
	}
	defer resp.Body.Close()

	// 读取响应
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("failed to read response: %v", err)
	}

	// 解析响应
	var chatResp ChatResponse
	if err := json.Unmarshal(respBody, &chatResp); err != nil {
		return "", fmt.Errorf("failed to parse response: %v", err)
	}

	if chatResp.Error != "" {
		return "", fmt.Errorf("API error: %s", chatResp.Error)
	}

	if len(chatResp.Choices) == 0 || chatResp.Choices[0].Message.Content == "" {
		return "", fmt.Errorf("no response from model")
	}

	// 应用代码高亮
	return formatCode(chatResp.Choices[0].Message.Content), nil
}

// interactiveMode 交互模式
func interactiveMode(model string, temperature float64, maxTokens int) {
	fmt.Println("### 软通动力: INTDOS-CodeMate 交互模式 ###")
	fmt.Println("### 进入交互模式 (输入 '/quit' 退出, '/clear' 清空屏幕):")
	scanner := bufio.NewScanner(os.Stdin)
	for {
		fmt.Print(">>> ")
		if !scanner.Scan() {
			fmt.Println("\n退出交互模式 <软通动力: INTDOS-CodeMate>")
			break
		}
		prompt := strings.TrimSpace(scanner.Text())
		if prompt == "" {
			continue
		}
		if prompt == "/quit" {
			fmt.Println("退出交互模式 <软通动力: INTDOS-CodeMate>")
			break
		}
		if prompt == "/clear" {
			fmt.Print("\033[H\033[2J") // 清屏
			continue
		}

		response, err := queryOllama(prompt, model, temperature, maxTokens)
		if err != nil {
			fmt.Fprintf(os.Stderr, "错误: %v\n", err)
			continue
		}
		fmt.Println(response)
	}
}

func main() {
	// 定义命令行参数
	prompt := flag.String("prompt", "", "单次查询的提示文本")
	interactive := flag.Bool("interactive", false, "进入交互模式")
	model := flag.String("model", defaultModel, "模型名称 (默认: "+defaultModel+")")
	temperature := flag.Float64("temperature", defaultTemperature, "生成温度 (默认: 0.7)")
	maxTokens := flag.Int("max-tokens", defaultMaxTokens, "最大生成令牌数 (默认: -1)")
	flag.Parse()

	// 检查参数
	if *interactive {
		interactiveMode(*model, *temperature, *maxTokens)
	} else if *prompt != "" {
		response, err := queryOllama(*prompt, *model, *temperature, *maxTokens)
		if err != nil {
			fmt.Fprintf(os.Stderr, "错误: %v\n", err)
			os.Exit(1)
		}
		fmt.Println(response)
	} else {
		flag.Usage()
		fmt.Fprintf(os.Stderr, "\n示例: ollama-cli -prompt \"写一个 Go HTTP 服务器\" 或 ollama-cli -interactive\n")
		os.Exit(1)
	}
}
